#include <iostream>
#include <stdio.h>

#include "arg_parser.h"
#include "lz4/lz4hc.h"

#include "zstd/zstd.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#include "rgb_to_nv12.h"

#define LOGE printf
#define LOGD printf

#define Y_BLOCK_WIDTH 16
#define UV_BLOCK_WIDTH 8

#define RESIDUAL_STEP 2


#define Y_BLOCK_SIZE (Y_BLOCK_WIDTH * Y_BLOCK_WIDTH)
#define UV_BLOCK_SIZE (UV_BLOCK_WIDTH * UV_BLOCK_WIDTH)

struct EncParam
{
    bool dump_frame = false;
    std::string ref_img;
    std::string cur_img;
    std::string next_img;

};

static void save_block_data2(uint8_t* out_frame, const uint8_t* src_frame, uint32_t block_x, uint32_t block_y, uint32_t width, uint32_t block_width)
{
    for (size_t i = 0; i < block_width; i++)
    {
        memcpy(out_frame + ((block_y + i) * width + block_x), src_frame + ((block_y + i) * width + block_x), block_width);        
    }
    
}

static void compare_block(const uint8_t* frame1, const uint8_t* frame2, uint32_t block_x, uint32_t block_y, uint32_t block_width, uint32_t width)
{
    for (size_t i = 0; i < block_width; i++)
    {
        for (size_t j = 0; j < block_width; j++)
        {
            uint32_t idx = (i + block_y) * width + block_x + j;
            if (frame1[idx] != frame2[idx])
            {
                LOGE("block not match at(x: %d, y: %d)\n", block_x, block_y);
                return ;
            }
        }        
    }    
}

static void block_search_nearest3(const uint8_t* reference_frame, const uint8_t* cur_frame, uint32_t width, uint32_t height, uint32_t block_width, uint32_t& pos_x, uint32_t& pos_y)
{
#if 1
    /**
     * @brief search rule: 
     *  #    #    #             #    #
     *  #    #    #             #    #
     *  #    #    #(cur block)  #    #
     *  #    #    #             #    # 
     *  #    #    #             #    # 
     * total compare block count: 25
     */
#define SEARCH_AREA (5 * block_width)

    uint32_t start_x = (pos_x > SEARCH_AREA) ? (pos_x - SEARCH_AREA) : 0;
    uint32_t start_y = (pos_y > SEARCH_AREA) ? (pos_y - SEARCH_AREA) : 0;

    uint32_t end_x = pos_x + SEARCH_AREA;
    uint32_t end_y = pos_y + SEARCH_AREA;

    end_x = (end_x > width) ? width : end_x;
    end_y = (end_y > height) ? height : end_y;
#else
    // full search
    uint32_t start_x = 0;
    uint32_t start_y = 0;

    uint32_t end_x = width;
    uint32_t end_y = height; 
#endif
    uint32_t out_x = 0;
    uint32_t out_y = 0;
    unsigned int min_diff = UINT_MAX;

    for (size_t i = start_y; i < end_y; i += block_width)
    {
        for (size_t j = start_x; j < end_x; j += block_width)
        {
            unsigned int diff_sum = 0;
            //const uint8_t* cur_buf = cur_frame + i * width + j;
            //const uint8_t* ref_buf = reference_frame + i * width + j;

            for (size_t block_h = 0; block_h < block_width; block_h++)
            {
                for (size_t block_w = 0; block_w < block_width; block_w++)
                {
                    unsigned int c1 = cur_frame[(pos_y + block_h) * width + pos_x + block_w];
                    unsigned int c2 = reference_frame[(i + block_h) * width + j + block_w];
                    diff_sum += (c1 - c2) * (c1 - c2);
                }                
            }

            if (min_diff > diff_sum)
            {
                out_x = j;
                out_y = i;
                min_diff = diff_sum;
            }
        }        
    }
#if 0
    if (min_diff != 0)
    {
        LOGD("[%s.%d]===>: beast diff: %u\n", __FUNCTION__, __LINE__, min_diff);
    }

    compare_block(reference_frame, cur_frame, out_x, out_y, block_width, width);
#endif
    pos_x = out_x;
    pos_y = out_y;
}

static void generate_predictive_frame(const uint8_t* ref_frame, const uint8_t* cur_frame, uint8_t* predictive_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            
            uint32_t pos_x = j;
            uint32_t pos_y = i;
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

            block_search_nearest3(ref_frame, cur_frame, width, height, block_width, pos_x, pos_y);
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            //printf("[%s.%d]===>: cur (%d,%d), nearest(%d,%d)\n", __FUNCTION__, __LINE__, j,i, pos_x, pos_y);

            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                memcpy(predictive_frame + (i + block_h) * width + j, ref_frame + (pos_y + block_h) * width + pos_x, block_width);
            }
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            
            //save_block_data2(predictive_frame, ref_frame, pos_x, pos_y, width, block_width);
        }
    }
}

static void partition_encode(const uint8_t* ref_frame, const uint8_t* enc_frame, uint16_t* coordinate_buf, uint8_t* residual_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    uint32_t block_count = 0;
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            
            uint32_t pos_x = j;
            uint32_t pos_y = i;
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

            block_search_nearest3(ref_frame, enc_frame, width, height, block_width, pos_x, pos_y);
            //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
            //printf("[%s.%d]===>: cur (%d,%d), nearest(%d,%d)\n", __FUNCTION__, __LINE__, j,i, pos_x, pos_y);
#if 0
            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                memcpy(predictive_frame + (i + block_h) * width + j, ref_frame + (pos_y + block_h) * width + pos_x, block_width);
                for (uint32_t block_w = 0; block_w < block_width; block_w++)
                {
                    /* code */
                }
                
            }
#else
            for (uint32_t block_h = 0; block_h < block_width; block_h++)
            {
                for (uint32_t block_w = 0; block_w < block_width; block_w++)
                {
                    uint32_t predictive_idx = (i + block_h) * width + j + block_w;
                    uint32_t ref_idx = (pos_y + block_h) * width + pos_x + block_w;
                    /** calculate residual */
                    residual_frame[predictive_idx] = (enc_frame[predictive_idx] - ref_frame[ref_idx]) / RESIDUAL_STEP; // + 0x80;
                }
            }

            block_count++;
            *coordinate_buf++ = pos_x;
            *coordinate_buf++ = pos_y;
#endif            
        }
    }

    LOGD("block_count: %d\n", block_count);
}

static void make_predictive_frame(const uint8_t* ref_frame, const uint16_t* coordinate_buf, uint8_t* predictive_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    //uint32_t block_size = block_width * block_width;
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            uint32_t pos_x = *coordinate_buf++;
            uint32_t pos_y = *coordinate_buf++;
            for (size_t block_h = 0; block_h < block_width; block_h++)
            {
                memcpy(predictive_frame + (i + block_h) * width + j, ref_frame + (pos_y + block_h)* width + pos_x, block_width);
            }            
        }   
    } 
}




static void dump_frame(const char* out_file, const uint8_t* frame_buf, uint32_t frame_size)
{
    FILE* fp = fopen(out_file, "wb");
    if (!fp)
    {
        LOGE("open file(%s) failed(%s)\n", out_file, strerror(errno));
        return ;
    }

    int res = fwrite(frame_buf, 1, frame_size, fp);

    fflush(fp);
    fclose(fp);
}


static void my_search(const uint8_t* reference_frame, const uint8_t* cur_frame, uint32_t frame_size, uint32_t start_idx, uint32_t block_width, uint32_t& out_idx)
{
#define MY_SEARCH_AREA (5 * block_width)    

    uint32_t start_pos = (start_idx > MY_SEARCH_AREA) ? (start_idx - MY_SEARCH_AREA) : 0;
    uint32_t end_pos = start_idx + MY_SEARCH_AREA;
    end_pos = (end_pos > frame_size) ? frame_size : end_pos;
    unsigned int min_diff = UINT_MAX;

    for (size_t i = start_idx; i < end_pos; i += block_width)
    {
        /* code */
        uint32_t sum_diff = 0;
        for (size_t j = 0; j < block_width; j++)
        {
            uint32_t c = reference_frame[i + j] - cur_frame[start_idx + j];
            sum_diff += c * c;
        }
        
        if (min_diff > sum_diff)
        {
            min_diff = sum_diff;
            out_idx = i;
        }
    }
    
}

static void my_encode(const uint8_t* ref_frame, const uint8_t* enc_frame, uint8_t* predictive_frame, uint8_t* residual_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    uint32_t frame_size = width * height;
    for (size_t i = 0; i < frame_size; i += block_width)
    {
        uint32_t pos = 0;
        my_search(ref_frame, enc_frame, frame_size, i, block_width, pos);
        memcpy(predictive_frame + i, ref_frame + pos, block_width);
        const uint8_t* p_cur = enc_frame + i;
        const uint8_t* p_predictive = predictive_frame + i;
        uint8_t* p_residual = residual_frame + i;
        for (size_t j = 0; j < block_width; j++)
        {
            p_residual[j] = (p_cur[j] - p_predictive[j]) / RESIDUAL_STEP;
        }        
    }    
}


static void test_enc(EncParam& param)
{
    uint32_t img_width = 0;
    uint32_t img_height = 0;
    int pic1_n;

    unsigned char* ref_frame = stbi_load(param.ref_img.c_str(), (int*)&img_width, (int*)&img_height, &pic1_n, 0);
    unsigned char* enc_frame = stbi_load(param.cur_img.c_str(), (int*)&img_width, (int*)&img_height, &pic1_n, 0);
    unsigned char* next_frame = stbi_load(param.next_img.c_str(), (int*)&img_width, (int*)&img_height, &pic1_n, 0);

    if (!ref_frame || !enc_frame || !next_frame)
    {
        if (ref_frame)
        {
            free(ref_frame);
        }
        else
        {
            LOGE("load ref frame(%s) failed\n", param.ref_img.c_str());
        }

        if (enc_frame)
        {
            free(enc_frame);
        }
        else
        {
            LOGE("load enc frame(%s)  failed\n", param.cur_img.c_str());
        }


        if (next_frame)
        {
            free(next_frame);
        }
        else
        {
            LOGE("load next frame(%s)  failed\n", param.next_img.c_str());
        }

        return ;
    }


    uint32_t y_size = img_width * img_height;
    uint32_t frame_size = y_size * 3 / 2;
    uint8_t* ref_buf = new uint8_t[frame_size];
    uint8_t* enc_buf = new uint8_t[frame_size];
    uint8_t* next_buf = new uint8_t[frame_size];

    rgb_to_nv12(ref_frame, pic1_n, img_width, img_height, ref_buf); 
    rgb_to_nv12(enc_frame, pic1_n, img_width, img_height, enc_buf);
    rgb_to_nv12(next_frame, pic1_n, img_width, img_height, next_buf);

    free(ref_frame);
    free(enc_frame);
    free(next_frame);

    uint32_t y_coordinate_size = y_size / Y_BLOCK_SIZE * 4; /** uint16_t for coordinate, so a pair (x,y) is 4 byte */
    uint32_t uv_coordinate_size = y_size / 2 / UV_BLOCK_SIZE * 4; /** uint16_t for coordinate, so a pair (x,y) is 4 byte */
    uint32_t coordinate_size = y_coordinate_size + uv_coordinate_size;

    uint8_t* predictive_frame = new uint8_t[frame_size];
    uint8_t* residual_buf = new uint8_t[coordinate_size + frame_size];
    uint8_t* dec_frame = new uint8_t[frame_size];

    uint8_t* residual_frame = residual_buf + coordinate_size;
    uint16_t* y_coordinate_buf = (uint16_t*)residual_buf;
    uint16_t* uv_coordinate_buf = (uint16_t*)(residual_buf + y_coordinate_size);

 #if 0   
    //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    generate_predictive_frame(ref_buf, enc_buf, predictive_frame, img_width, img_height, Y_BLOCK_WIDTH);

    

    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    /** calculate residual */
    for (uint32_t i = 0; i < y_size; i++)
    {
        residual_frame[i] = (enc_buf[i] - predictive_frame[i]) / 2 + 0x80;
    }
#else
#if 1
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    partition_encode(ref_buf, enc_buf, y_coordinate_buf, residual_frame, img_width, img_height, Y_BLOCK_WIDTH);
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    partition_encode(ref_buf + y_size, enc_buf + y_size, uv_coordinate_buf, residual_frame + y_size, img_width / 2, img_height, UV_BLOCK_WIDTH);
#else
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    my_encode(ref_buf, enc_buf, predictive_frame, residual_frame, img_width, img_height, Y_BLOCK_SIZE);
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    my_encode(ref_buf + y_size, enc_buf + y_size, predictive_frame + y_size, residual_frame + y_size, img_width / 2, img_height, UV_BLOCK_SIZE);
#endif
#endif
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

    if (param.dump_frame)
    {
        //decode frame

        LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
        make_predictive_frame(ref_buf, y_coordinate_buf, predictive_frame, img_width, img_height, Y_BLOCK_WIDTH);
        LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
        make_predictive_frame(ref_buf + y_size, uv_coordinate_buf, predictive_frame + y_size, img_width / 2, img_height, UV_BLOCK_WIDTH);
        LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

        for (uint32_t i = 0; i < frame_size; i++)
        {
            dec_frame[i] = predictive_frame[i] + ((residual_frame[i] /*- 0x80*/) * RESIDUAL_STEP);
        }
        
        dump_frame("enc_frame.nv12", enc_buf, frame_size);
        dump_frame("ref_frame.nv12", ref_buf, frame_size);
        dump_frame("predictive_frame.nv12", predictive_frame, frame_size);
        dump_frame("residual_frame.nv12", residual_frame, frame_size);
        dump_frame("dec_frame.nv12", dec_frame, frame_size);        
    }

    uint8_t* temp = ref_buf;

    ref_buf = dec_frame; //enc_buf; // or predictive_frame    
    dec_frame = temp;

    temp = enc_buf;
    enc_buf = next_buf;
    next_buf = temp;

    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    partition_encode(ref_buf, enc_buf, y_coordinate_buf, residual_frame, img_width, img_height, Y_BLOCK_WIDTH);
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    partition_encode(ref_buf + y_size, enc_buf + y_size, uv_coordinate_buf, residual_frame + y_size, img_width / 2, img_height, UV_BLOCK_WIDTH);

    if (param.dump_frame)
    {
        make_predictive_frame(ref_buf, y_coordinate_buf, predictive_frame, img_width, img_height, Y_BLOCK_WIDTH);
        make_predictive_frame(ref_buf + y_size, uv_coordinate_buf, predictive_frame + y_size, img_width / 2, img_height, UV_BLOCK_WIDTH);
        //decode frame
        for (uint32_t i = 0; i < frame_size; i++)
        {
            dec_frame[i] = predictive_frame[i] + ((residual_frame[i] /*- 0x80*/) * RESIDUAL_STEP);
        }

        dump_frame("next_frame.nv12", enc_buf, frame_size);
        //dump_frame("ref_frame.nv12", ref_buf, frame_size);
        dump_frame("next_predictive_frame.nv12", predictive_frame, frame_size);
        dump_frame("next_residual_frame.nv12", residual_frame, frame_size);
        dump_frame("next_dec_frame.nv12", dec_frame, frame_size);        
    }


#if 1
    uint32_t out_buf_size = LZ4_COMPRESSBOUND(frame_size);
    uint8_t* out_buf = new uint8_t[out_buf_size];
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);

    int res_size = LZ4_compress_HC((const char*)residual_buf, (char*)out_buf, frame_size, out_buf_size, LZ4HC_CLEVEL_MAX);

    LOGD("compress residual: before: %d, after: %d\n", frame_size, res_size);

    res_size = LZ4_compress_HC((const char*)enc_buf, (char*)out_buf, frame_size, out_buf_size, LZ4HC_CLEVEL_MAX);

    LOGD("compress org: before: %d, after: %d\n", frame_size, res_size);
#else  
    uint32_t out_buf_size = ZSTD_COMPRESSBOUND(frame_size);
    uint8_t* out_buf = new uint8_t[out_buf_size];
    LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    int res_size = ZSTD_compress((char*)out_buf, out_buf_size, (const char*)residual_frame, frame_size, 22);

    LOGD("compress residual: before: %d, after: %d\n", frame_size, res_size);

    res_size = ZSTD_compress((char*)out_buf, out_buf_size, (const char*)enc_buf, frame_size, 22);

    LOGD("compress org: before: %d, after: %d\n", frame_size, res_size);
#endif

    delete[] predictive_frame;
    delete[] enc_buf;
    delete[] ref_buf;
    delete[] residual_buf;
    delete[] out_buf;

}


int main(int argc, char* argv[])
{
    EncParam maker_param;

    icpp::core::ArgParser a;

    a.addWithParamNoneHandler("help",'h', "show this usage info", [&a]() -> icpp::core::ArgumentParseError
    {
        a.showHelp();
        return icpp::core::ArgumentParseError::kDoCommand;
    });

    a.addWithParamRequired<std::string>("reference", 'r', "reference image file", true, maker_param.ref_img);
    a.addWithParamRequired<std::string>("encode", 'e', "the image will be encode", true, maker_param.cur_img);
    a.addWithParamRequired<std::string>("next", 'n', "next image will be encode", true, maker_param.next_img);
    a.addWithParamNone("dump", 'd', "is need dump processed frame", maker_param.dump_frame);

    icpp::core::ArgumentParseError ret = a.onParse(argc, argv);

    if (icpp::core::ArgumentParseError::kSuccess == ret)
    {
        test_enc(maker_param);       
    }

    return 0;

}

